//
//  MAPYShadowSocksManager.m
//  Mars Proxy
//
//  Created by PowerfullClean on 2023/10/17.
//

#import "MAPYShadowSocksManager.h"
#import "MAPYShadowSocksHandler.h"

#include <arpa/inet.h>
#include <limits.h>
#include <pthread.h>


#import <Shadowsocks_iOS/Shadowsocks_iOS.h>
#import <PacketProcessor_iOS/PacketProcessor_iOS.h>


const int kShadowSocksLocalPort = 9999;
static const int kShadowSocksTimeoutSecs = 600;
static const int kShadowSocksTcpAndUdpMode = 1;
static char *const kShadowSocksLocalAddress = "127.0.0.1";

static int get_port_of_socket (int fd) {
    struct sockaddr_in sin;
    socklen_t len = sizeof(sin);
    if (getsockname(fd, (struct sockaddr *)&sin, &len) < 0) {
        return 0;
    } else {
        return ntohs(sin.sin_port);
    }
}

@interface MAPYShadowSocksManager ()

@property (nonatomic) bool checkConnectivity;
@property (nonatomic) pthread_t ssLocalThreadId;

@property (nonatomic, copy) void (^startCompletion)(ErrorCode);
@property (nonatomic, copy) void (^stopCompletion)(ErrorCode);

@property (nonatomic) dispatch_queue_t dispatchQueue;
@property (nonatomic) dispatch_group_t dispatchGroup;

@property (nonatomic) MAPYShadowSocksHandler *ssConnectivity;

@end


@implementation MAPYShadowSocksManager

- (id)init:(NSDictionary *)config {
  self = [super init];
  if (self) {
    _config = config;
    _dispatchQueue = dispatch_queue_create("OneShadowSocksManager", DISPATCH_QUEUE_SERIAL);
    _dispatchGroup = dispatch_group_create();
    _ssConnectivity = [[MAPYShadowSocksHandler alloc] initWithPort:kShadowSocksLocalPort];
  }
  return self;
}

- (void)startWithConnectivityChecks:(bool)checkConnectivity
                         completion:(void (^)(ErrorCode))completion {
  if (self.ssLocalThreadId != 0) {
    return completion(ShadowSocksStartFailure);
  }
  self.checkConnectivity = checkConnectivity;
  dispatch_async(dispatch_get_main_queue(), ^{
    // Start ss-local from the main application thread.
    self.startCompletion = completion;
    [self startShadowSocksThread];
  });
}

-(void)stop:(void (^)(ErrorCode))completion {
  if (self.ssLocalThreadId == 0) {
    return;
  }
  dispatch_async(dispatch_get_main_queue(), ^{
    self.stopCompletion = completion;
    pthread_kill(self.ssLocalThreadId, SIGUSR1);
    self.ssLocalThreadId = 0;
  });
}

#pragma mark - Lifecycle

void ShadowSocksCallback(int socks_fd, int udp_fd, void *udata) {
  DLog(@"OneShadowSocksManager callback.");
  if (socks_fd <= 0 || udp_fd <= 0) {
    return;
  }
  MAPYShadowSocksManager* ss = (__bridge MAPYShadowSocksManager *)udata;
  ss.localSSPort = get_port_of_socket(socks_fd);
  ss.localSSUdpPort = get_port_of_socket(udp_fd);
  [ss checkServerConnectivity];
}

- (void)startShadowSocks {
  if (self.config == nil) {
    self.startCompletion(illegalServerConfiguration);
     return;
  }
  int port = [self.config[@"port"] intValue];
  char *host = (char *)[self.config[@"host"] UTF8String];
  char *password = (char *)[self.config[@"password"] UTF8String];
  char *method = (char *)[self.config[@"method"] UTF8String];
  const profile_t profile = {.remote_host = host,
                             .local_addr = kShadowSocksLocalAddress,
                             .method = method,
                             .password = password,
                             .remote_port = port,
                             .local_port = kShadowSocksLocalPort,
                             .timeout = kShadowSocksTimeoutSecs,
                             .acl = NULL,
                             .log = NULL,
                             .fast_open = 0,
                             .mode = kShadowSocksTcpAndUdpMode,
                             .verbose = 0};
    int success = start_ss_local_server_with_callback(profile, ShadowSocksCallback, (__bridge void *)self);
    if (success < 0) {
      self.startCompletion(ShadowSocksStartFailure);
      return;
    }
 
    if (self.stopCompletion) {
        self.stopCompletion(noError);
        self.stopCompletion = nil;
    }
}

// Entry point for the OneShadowSocksManager POSIX thread.
void *startShadowSocks(void *udata) {
  MAPYShadowSocksManager* ss = (__bridge MAPYShadowSocksManager *)udata;
  [ss startShadowSocks];
  return NULL;
}

// Starts a POSIX thread that runs ss-local.
- (void)startShadowSocksThread {
  pthread_attr_t attr;
  int err = pthread_attr_init(&attr);
  if (err) {
     self.startCompletion(ShadowSocksStartFailure);
    return;
  }
  err = pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);
  if (err) {
     self.startCompletion(ShadowSocksStartFailure);
    return;
  }
  err = pthread_create(&_ssLocalThreadId, &attr, &startShadowSocks, (__bridge void *)self);
  if (err) {
     self.startCompletion(ShadowSocksStartFailure);
  }
  err = pthread_attr_destroy(&attr);
  if (err) {
     self.startCompletion(ShadowSocksStartFailure);
    return;
  }
}

#pragma mark - Connectivity
- (void)checkServerConnectivity {
  if (!self.checkConnectivity) {
    self.startCompletion(noError);
    return;
  }
  __block BOOL isRemoteUdpShadowSocksAvailabled = false;
  __block BOOL serverCredentialsAreValid = false;
  __block BOOL isServerReachable = false;

  // Enter the group once for each check
  dispatch_group_enter(self.dispatchGroup);
  dispatch_group_enter(self.dispatchGroup);
  dispatch_group_enter(self.dispatchGroup);

  [self.ssConnectivity isUdpShadowSocksAvailabled:^(BOOL enabled) {
    isRemoteUdpShadowSocksAvailabled = enabled;
    dispatch_group_leave(self.dispatchGroup);
  }];
  [self.ssConnectivity isReachable:self.config[@"host"]
                              port:[self.config[@"port"] intValue]
                        completion:^(BOOL isReachable) {
                          isServerReachable = isReachable;
                          dispatch_group_leave(self.dispatchGroup);
                        }];
  [self.ssConnectivity checkServerCredentials:^(BOOL valid) {
    serverCredentialsAreValid = valid;
    dispatch_group_leave(self.dispatchGroup);
  }];

  dispatch_group_notify(self.dispatchGroup, self.dispatchQueue, ^{
    DLog(@"Server connectivity checks done");
    if (isRemoteUdpShadowSocksAvailabled) {
      self.startCompletion(noError);
    } else if (serverCredentialsAreValid) {
      self.startCompletion(udpRelayNotEnabled);
    } else if (isServerReachable) {
      self.startCompletion(invalidServerCredentials);
    } else {
      self.startCompletion(serverUnreachable);
    }
  });
}

- (void)isReachable:(void (^)(ErrorCode))completion {
    
    __block BOOL isServerReachable = false;
    
    // Enter the group once for each check
    dispatch_group_enter(self.dispatchGroup);
    
    [self.ssConnectivity isReachable:self.config[@"host"]
                                port:[self.config[@"port"] intValue]
                          completion:^(BOOL isReachable) {
        isServerReachable = isReachable;
        dispatch_group_leave(self.dispatchGroup);
    }];
    
    dispatch_group_notify(self.dispatchGroup, self.dispatchQueue, ^{
        DLog(@"Server connectivity checks done");
        completion(isServerReachable ? noError : serverUnreachable);
    });
}

- (BOOL)startCheckConnectivity:(void (^)(BOOL))completion {
    
    __block BOOL serverConnectivityCheck = false;
    
    // Enter the group once for each check
    dispatch_group_enter(self.dispatchGroup);

    [self.ssConnectivity checkServerConnectivity:^(BOOL check) {
        serverConnectivityCheck = check;
        dispatch_group_leave(self.dispatchGroup);
    }];
    
    dispatch_group_notify(self.dispatchGroup, self.dispatchQueue, ^{
        DLog(@"Server connectivity checks done");
        completion(serverConnectivityCheck);
    });
    return YES;
}

@end

